package com.sparrow.common.hotswap;

import com.sparrow.common.util.ClassUtil;
import com.sparrow.common.util.CollectionUtil;
import com.sun.tools.attach.VirtualMachine;
import org.objectweb.asm.ClassReader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.lang.instrument.ClassDefinition;
import java.lang.instrument.Instrumentation;
import java.lang.management.ManagementFactory;
import java.lang.reflect.Field;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.HashMap;
import java.util.Map;

/**
 * 用于运行时重新定义已经加载的类
 */
public class HotSwapManager {

    private static Logger logger = LoggerFactory.getLogger(HotSwapManager.class);

    private static Instrumentation instrumentation;

    private final Map<Path, Long> lastModifiedCache = new HashMap<>();

    private String hotswapClassDir;

    private static String agentPath;

    public HotSwapManager(String hotswapClassDir, String agentPath) {
        this.hotswapClassDir = hotswapClassDir;
        this.agentPath = agentPath;
    }

    /**
     * 尝试重定义class，有修改的class才会重定义
     */
    public void tryRedefineClass() {
        // 获取指定目录下所有的.class文件
        Path dir = Paths.get(hotswapClassDir);
        Map<String, byte[]> classByteMap = new HashMap<>();
        try {
            Files.walk(dir)
                    .filter(path -> path.toString().endsWith(".class"))
                    .forEach(path -> {
                        try {
                            long currentModified = Files.getLastModifiedTime(path).toMillis();
                            Long cachedModified = lastModifiedCache.get(path);

                            if (cachedModified == null || currentModified > cachedModified) {
                                byte[] classBytes = Files.readAllBytes(path);
                                String className = ClassUtil.getClassName(classBytes);
                                classByteMap.put(className, classBytes);

                                // 更新缓存时间戳
                                lastModifiedCache.put(path, currentModified);
                            }
                        } catch (IOException e) {
                            logger.error("读取class文件失败：{}", path, e);
                        }
                    });
        } catch (IOException e) {
            logger.error("热更错误", e);
        }
        if (CollectionUtil.isEmpty(classByteMap)) {
            return;
        }

        ClassDefinition[] classDefinitions = new ClassDefinition[classByteMap.size()];
        int index = 0;
        for (Map.Entry<String, byte[]> e : classByteMap.entrySet()) {
            try {
                Class<?> clazz = Class.forName(e.getKey());
                classDefinitions[index++] = new ClassDefinition(clazz, e.getValue());
            } catch (Exception ee) {
                throw new RuntimeException("构造 ClassDefinition 失败，class: " + e.getKey(), ee);
            }
        }
        // 重新定义所有的类
        try {
            Instrumentation inst = getInstrumentation();
            inst.redefineClasses(classDefinitions);
        } catch (Exception e) {
            throw new RuntimeException("重新定义类失败", e);
        }

        for (Map.Entry<String, byte[]> e : classByteMap.entrySet()) {
            logger.info(String.format("热更成功 %s.class", e.getKey()));
        }
    }

    private static void loadAgent(String agentPath) {
        VirtualMachine vm = null;
        try {
            String processName = ManagementFactory.getRuntimeMXBean().getName();
            String pid = processName.split("@")[0];
            vm = VirtualMachine.attach(pid);
            vm.loadAgent(agentPath);
        } catch (Exception ex) {
            throw new RuntimeException(ex);
        } finally {
            if (vm != null) {
                try {
                    vm.detach();
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
        }
    }

    /**
     * 通过反射初始化设置instrumentation引用
     * @return
     */
    public static synchronized Instrumentation getInstrumentation() {
        if (instrumentation == null) {
            loadAgent(agentPath);
            try {
                Class<?> clazz = Class.forName("com.agent.Premain");
                Field field = clazz.getDeclaredField("instrumentation");
                field.setAccessible(true);
                instrumentation = (Instrumentation) field.get(null);
            } catch (Exception e) {
                throw new RuntimeException("加载 Agent 失败", e);
            }
        }
        return instrumentation;
    }
}
